(*
 * Store of management
 *
 * currently a bit inefficient; the management->id mapping should be
 * some kind of trie ultimately 
 *)

open ExtList
open Printf

open Management;;

module Int =
struct
  type t = int
  let compare = compare
end ;;


module Float =
struct
  type t = float
  let compare = compare
end ;;

module IdMap = Map.Make (Int);;

module MgtMap = Map.Make (Management) ;;

type mgtID = int

type score = float ;;

type statistics = {
    min_sc : score;
    max_sc : score;
    mean_sc : score;
    derived_sc : score;
    attempts : int;
    insolvencies : int }

       
let derive_sc stat = {
  stat with derived_sc = 
		      match stat.attempts with 
		      | 0 | 1 | 2 | 3 -> stat.mean_sc
		      | _ when stat.attempts < 10 -> stat.mean_sc *. 0.9
		      | _ -> stat.mean_sc *. 0.1 }


let newst sc = derive_sc {
  min_sc = sc;
  max_sc = sc;
  mean_sc = sc;
  derived_sc = sc;
  attempts = 1;
  insolvencies = 
  if sc > 0.0
  then 0
  else 1 }
		   

let updateStats s sc = derive_sc {
  min_sc = if sc < s.min_sc then sc else s.min_sc;
  max_sc = if sc > s.max_sc then sc else s.max_sc;
  mean_sc = (s.mean_sc *. (float_of_int s.attempts) +. sc) /. float_of_int (s.attempts + 1);
  attempts = s.attempts + 1;
  derived_sc = 0.;
  insolvencies = s.insolvencies + 
  if sc > 0.0
  then 0 
  else 1 }
    

type managementStore = {
    ms_size : int;
    mgt_to_id: (mgtID) MgtMap.t;
    id_to_st: (statistics) IdMap.t;
    id_to_mgt: management IdMap.t;
    sc_to_id: (score * mgtID) list
  }


let assoc_del_val v l =
  List.filter (fun (_,x) -> x != v) l 

let rec assoc_ins k v l = 
  match l with
  | [] -> [(k,v)]
  | ([(k1,v1)]) as al when k >= k1 -> (k,v) :: al
  | [(k1,v1)] when k >= k1 -> [(k1,v1);(k,v)] 
  | [(k1,v1) as al]  -> [al;(k,v)] 
      (* I think for NaNs - certainly bugged here *)
  | ((k1,v1)::((k2,v2)::l) ) as ar when k >= k1 -> (k,v) :: ar 
  | ((k1,v1)::(((k2,v2)::l) as r )) when k < k1 && k >= k2 -> [(k1,v1);(k,v)] @ r
  | ((k1,v1)::(((k2,v2)::l) as r )) -> (k1,v1) :: (assoc_ins k v r)

let insertNew ms mgt sc =
  let nid =  (ms.ms_size + 1) in
  { (* ms with  *)
    ms_size = nid;
    id_to_st = IdMap.add nid (newst sc) ms.id_to_st;
    mgt_to_id = MgtMap.add mgt nid ms.mgt_to_id ;
    id_to_mgt = IdMap.add nid mgt ms.id_to_mgt ;
    sc_to_id = assoc_ins sc nid ms.sc_to_id }

let updateExists ms mgt sc =
  let id = MgtMap.find mgt ms.mgt_to_id in
  let stats = IdMap.find id ms.id_to_st in
  let nstats = updateStats stats sc in
  { ms with
    sc_to_id = assoc_ins (nstats.derived_sc) id (assoc_del_val id ms.sc_to_id);
    id_to_st = IdMap.add id nstats (IdMap.remove id ms.id_to_st) }

let updateEntry ms mgt sc =
  if MgtMap.mem mgt ms.mgt_to_id 
  then updateExists ms mgt sc
  else insertNew ms mgt sc

let getBestMgt n ms = 
  let ids = List.map snd (List.take n ms.sc_to_id) in
  List.map (fun i -> IdMap.find i ms.id_to_mgt) ids

let mlkp ms mid =
  IdMap.find mid ms.id_to_mgt

let print_table_best ms =
  match ms.sc_to_id with 
  | [] -> ""
  | _ -> (print_management (mlkp ms (snd (List.hd ms.sc_to_id))))

let print_managementstore ms =
  let (scs,ids) = (List.split ms.sc_to_id) in
  let mgts = List.map (fun id -> (print_management (mlkp ms id))) ids in
  List.fold_left2 (fun a b c -> a ^ b ^ " : " ^ c ) "" (List.map (sprintf "%8.3f") scs) mgts

let addList ms l = 
  List.fold_left (fun ms0 (sc,mgt) -> updateEntry ms0 mgt sc) ms l


let empty_management_store =
  { ms_size = 0 ;
    mgt_to_id = MgtMap.empty ;
    id_to_st = IdMap.empty ;
    id_to_mgt = IdMap.empty ;
    sc_to_id = [] }
